﻿import os
import uvicorn
from fastapi import FastAPI, File, UploadFile
from starlette.middleware.cors import CORSMiddleware
from config import UPLOAD_PATH
from logs import LOGGER
from milvus_helpers import MilvusHelper
from encode import SentenceModel
from operations.load import do_load,adddata
from operations.search import do_search
from operations.count import do_count
from operations.drop import do_drop
import json
import numpy
from pydantic import BaseModel

class EncodeParam(BaseModel):
    content: str
    tableName:str=None
    embedding:str=None

class NumpyEncoder(json.JSONEncoder):  
    def default(self, obj):  
        if isinstance(obj, (numpy.int_, numpy.intc, numpy.intp, numpy.int8,  
            numpy.int16, numpy.int32, numpy.int64, numpy.uint8,  
            numpy.uint16, numpy.uint32, numpy.uint64)):  
            return int(obj)  
        elif isinstance(obj, (numpy.float_, numpy.float16, numpy.float32,numpy.float64)):  
            return float(obj)  
        elif isinstance(obj, (numpy.ndarray,)):  
            return obj.tolist()  
        return json.JSONEncoder.default(self, obj) 

app = FastAPI()
origins = ["*"]
app.add_middleware(
    CORSMiddleware,

    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],

)

MODEL = SentenceModel()
MILVUS_CLI = MilvusHelper()
# Mkdir '/tmp/qa-data'  
if not os.path.exists(UPLOAD_PATH):
    os.makedirs(UPLOAD_PATH)


@app.post('/qa/load_data')
async def do_load_api(file: UploadFile = File(...), table_name: str = None):
    try:
        text = await file.read()
        fname_path = os.path.join(UPLOAD_PATH, file.filename)
        with open(fname_path, 'wb') as f:
            f.write(text)
        total_num = do_load(table_name, fname_path, MODEL, MILVUS_CLI)
        LOGGER.info(f"Successfully loaded data, total count: {total_num}")
        return {'status': True, 'msg': f"Successfully loaded data: {total_num}"}, 200
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400


@app.post('/qa/search')
async def do_get_question_api(param: EncodeParam):
    try:
        questions, distances= do_search(param.tableName, param.content, MODEL, MILVUS_CLI,param.embedding)
        return {'code': "1", 'body': (questions,distances)}
    except Exception as e:
        LOGGER.error(e)
        return {'code': "0", 'message': ""}


@app.post('/qa/count')
async def count_images(table_name: str = None):
    try:
        num = do_count(table_name, MILVUS_CLI)
        LOGGER.info("Successfully count the number of questions!")
        return num
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400


@app.post('/qa/drop')
async def drop_tables(table_name: str = None):
    try:
        status = do_drop(table_name, MILVUS_CLI)
        return {'status': True, 'msg': status}
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400

@app.post('/qa/encode')
async def encodecontent(param: EncodeParam):
    try:
        em= MODEL.sentence_encode([param.content])
        return {'code': "1", 'message': json.dumps(em[0],cls=NumpyEncoder,ensure_ascii=False)}
    except Exception as e:
        LOGGER.error(e)
        return {'code': "0", 'message': ""}

@app.get('/qa/adddata')
async def do_load_api(filename: str = None,table_name: str = None):
    try:
        total_num = adddata(table_name, filename, MODEL, MILVUS_CLI)
        return {'code': "1", 'message': total_num}
    except Exception as e:
        LOGGER.error(e)
        return {'code': "0", 'message': ""}

if __name__ == '__main__':
    uvicorn.run(app=app, host='0.0.0.0', port=8000)
